/*-------------------------------------------------------------------------
 *
 * txn_timestmap.c
 *  timestamp coordination of transactions
 *
 * Copyright (c) 2020, Alibaba Group Holding Limited
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *   http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. *
 * src/backend/distributed_txn/txn_timestamp.c
 *
 *-------------------------------------------------------------------------
 */

#include "postgres.h"

#include "access/xact.h"
#include "access/transam.h"
#include "access/twophase.h"
#include "distributed_txn/txn_timestamp.h"
#include "fmgr.h"
#include "libpq/libpq.h"
#include "utils/builtins.h"
#include "access/mvccvars.h"
#include "storage/proc.h"
#include "storage/procarray.h"

/* user-set guc parameter */
bool		enable_timestamp_debug_print = false;

/*
 * start timestamp used in both coordinator and worker:
 */
static bool backendReceivedTimestamp = false;

/* start_ts is generated by node itself or coordinated */
static LogicalTime txnStartTs = 0;
/* commit_ts = ClockTick() */
static LogicalTime txnCommitTs = 0;
/* coordinated_ts = Max(prepare_ts) */
static LogicalTime txnCoordinatedCommitTs = 0;
static LogicalTime replyTs = 0;

IsCoordinating2PCHook IsCoordinating2PC = NULL;

static void TxnSetStartTs(LogicalTime);
static void TxnSetCommitTs(LogicalTime);

LogicalTime
TxnGetStartTs(void)
{
	return txnStartTs;
}

LogicalTime
TxnGetCoordinatedCommitTs(void)
{
	if (enable_timestamp_debug_print)
		elog(LOG, "Get coordinated committs " UINT64_FORMAT " to send to workers", txnCoordinatedCommitTs);

	return txnCoordinatedCommitTs;
}

/*
 * Called by GetSnapshotData on Coordinator or Worker
 * On coordinator:
 *      use ClockNow(), since GetSnapshotData is determined by XactIsoLevel
 * On worker:
 *      in RC/RR,  receiving start_ts happens before GetSnapshotData
 * We rule them all in following ways:
 * 1. TxnGenerateStartTs() initializes txnStartTs only at the transaction start
 * 2. BackendRecvTimestamp() updates txnStartTs when receving timestamp from coordinator,
 *      so the next time GetSnapshotData() will return the latest txnStartTs
 */
LogicalTime
TxnGetOrGenerateStartTs(bool latest)
{
	LogicalTime res;

	if (!txnUseGlobalSnapshot || latest)
	{
		res = LogicalClockNow();
		if (enable_timestamp_debug_print)
		{
			elog(LOG, "PhysicalNow: " LOGICALTIME_FORMAT,
				 LOGICALTIME_STRING(PhysicalClockNow()));
			const char *snapshotFreshness = latest ? "latest" : "stable";
			const char *snapshotCreator = txnUseGlobalSnapshot ? "global" : "local";

			elog(LOG, "Generate start_ts " UINT64_FORMAT LOGICALTIME_FORMAT
				 " use %s %s snapshot MyTmin " UINT64_FORMAT " procno %d",
				 res, LOGICALTIME_STRING(res), snapshotFreshness, snapshotCreator,
				 pg_atomic_read_u64(&MyPgXact->tmin), MyProc->pgprocno);
		}
		if (!latest)
			txnStartTs = res;
	}
	else
	{
		res = txnStartTs;
		if (enable_timestamp_debug_print)
		{
			const char *snapshotCreator = txnUseGlobalSnapshot ? "global" : "local";

			elog(LOG, "Get start_ts " UINT64_FORMAT LOGICALTIME_FORMAT " from coordinator "
				 "using %s snapshot MyTmin " UINT64_FORMAT " procno %d",
				 res, LOGICALTIME_STRING(res), snapshotCreator, pg_atomic_read_u64(&MyPgXact->tmin), MyProc->pgprocno);
		}
	}
	return res;
}

/**
 * Called by TransactionIdAsyncCommitTree() on Coordinator or Worker
 * So it needs to meet the requirements of 2PC, 1PC, 0PC on both coordinator and worker.
 *
 * On coordinator:
 *      use coordinated_commit_ts if in 2PC transaction, else use ClockTick()
 * On worker:
 *      use commit_ts assigned from txn_commit_prepare() if in 2PC transaction,
 *      else use ClockTick()
 */
LogicalTime
TxnGetOrGenerateCommitTs(bool fromCoordinator)
{
	LogicalTime res;

	if (txnCoordinatedCommitTs != 0)
	{
		if (txnUseGlobalSnapshot && false == fromCoordinator)
			elog(ERROR, "Coordinator did not pass any commit timestamp down");

		/* coordinated ts from 2PC */
		res = txnCoordinatedCommitTs;
		if (enable_timestamp_debug_print)
			elog(LOG, "Get commit timestamp from coordinator " UINT64_FORMAT " logical clock " UINT64_FORMAT,
				 res, ToLogicalClock(res));
	}
	else
	{
		if (txnUseGlobalSnapshot && fromCoordinator)
			elog(ERROR, "Coordinator should not pass commit timestamp down for 1PC transactions");

		/* generate commit_ts by myself */
		if (txnCommitTs)
			elog(ERROR, "commit timestamp exists unexpectly " UINT64_FORMAT, txnCommitTs);

		TxnSetCommitTs(LogicalClockTick());
		res = txnCommitTs;
		if (enable_timestamp_debug_print)
			elog(LOG, "Get commit timestamp locally " UINT64_FORMAT " logical clock " UINT64_FORMAT,
				 res, ToLogicalClock(res));
	}
	Assert(!COMMITSEQNO_IS_SUBTRANS(res));
	return res;
}

void
AtEOXact_txn(void)
{
	if (enable_timestamp_debug_print)
		elog(LOG, "AtEOXact_txn");

	txnUseGlobalSnapshot = false;
	RecentGlobalTs = InvalidCommitSeqNo;
	txnCoordinatedCommitTs = 0;
	TxnSetStartTs(0);
	TxnSetCommitTs(0);
}

static void
TxnSetStartTs(LogicalTime startTs)
{
	txnStartTs = startTs;
	if (startTs)
		ereport(DEBUG1, (errmsg("set start_ts %lu", startTs)));
}

static void
TxnSetCommitTs(LogicalTime commitTs)
{
	txnCommitTs = commitTs;

	if (enable_timestamp_debug_print)
		elog(LOG, "set commit timestamp " UINT64_FORMAT, commitTs);

}

void
TxnSetCoordinatedCommitTs(LogicalTime ts)
{
	Assert(txnCoordinatedCommitTs == 0 || ts == 0);
	LogicalClockUpdate(ts);
	txnCoordinatedCommitTs = ts;

	if (enable_timestamp_debug_print)
		elog(LOG, "set commit timestamp from coordinator " UINT64_FORMAT, ts);

}

void
TxnSetCoordinatedCommitTsFromStr(const char *ts)
{
	LogicalTime commit_ts = (LogicalTime) strtoull(ts, NULL, 0);

	if (commit_ts == 0)
		ereport(ERROR, (errmsg("coordinated commit_ts is 0")));

	TxnSetCoordinatedCommitTs(commit_ts);
}

void
TxnSetReplyTimestamp(LogicalTime ts)
{
	replyTs = ts;
}

LogicalTime
TxnGetAndClearReplyTimestamp(void)
{
	LogicalTime res = replyTs;

	if (!backendReceivedTimestamp)
	{
		return 0;
	}
	replyTs = 0;

	if (enable_timestamp_debug_print)
		elog(LOG, "reply timestamp " UINT64_FORMAT " logical clock " UINT64_FORMAT,
			 res, ToLogicalClock(res));

	return res;
}

/*
 * Receive timestamp from coordinator
 * read-committed: receive start_ts every statement, use start_ts as snapshot.csn
 * repeatable-read: receive start_ts before first statement
 * 0PC: no explicit transaction exists, just update the clock but not set snapshot
*/
void
BackendRecvTimestamp(LogicalTime ts)
{
	/*
	 * Be carefull of the order between setting MyPgXact->tmin and acquiring
	 * the ts_lock to fetch maxCommitTs which is critical to the correctness
	 * of garbage collection algorithm. Written by Junbin Kang, 2020.01.20
	 */
	pg_atomic_write_u64(&MyPgXact->tmin, ts);
	pg_memory_barrier();

	backendReceivedTimestamp = true;
	RecentGlobalTs = LogicalClockUpdate(ts);
	TxnSetStartTs(ts);
	txnUseGlobalSnapshot = true;
	if (enable_timestamp_debug_print)
		elog(LOG, "recv start timestamp trans %d " UINT64_FORMAT " procno %d", IsTransactionState(), ts, MyProc->pgprocno);
}

/*
 * This timestamp from worker could be either prepare_ts or commit_ts
 * prepare_ts: coordinator in COORD_TRANS_STARTED state and sended prepare command
 * commit_ts: any other scenarios
*/
void
FrontendRecvTimestamp(LogicalTime ts)
{
	LogicalClockUpdate(ts);
	if (enable_timestamp_debug_print)
		elog(LOG, "recv reply timestamp trans " UINT64_FORMAT " logical clock " UINT64_FORMAT,
			 ts, ToLogicalClock(ts));
}

PG_FUNCTION_INFO_V1(txn_get_start_ts);
PG_FUNCTION_INFO_V1(txn_commit_prepared);

/**
 * inspect start_ts of a transaction, for test purpose
 */
Datum
txn_get_start_ts(PG_FUNCTION_ARGS)
{
	LogicalTime ts = TxnGetStartTs();

	PG_RETURN_UINT64(ts);
}

/**
 * Commit a prepared transaction, with timestamp assigned from coordinator
 */
Datum
txn_commit_prepared(PG_FUNCTION_ARGS)
{
	const char *gid = text_to_cstring(PG_GETARG_TEXT_PP(0));
	LogicalTime commit_ts = PG_GETARG_INT64(1);

	if (commit_ts == 0)
		ereport(ERROR, (errmsg("coordinated commit_ts is 0")));

	TxnSetCoordinatedCommitTs(commit_ts);
	FinishPreparedTransaction(gid, true);
	PG_RETURN_NULL();
}
